perf(moe): add frozen grouped GEMM for adapter training - #3

Written by Codex.

Why this change

Nemotron-3 Super LoRA freezes the 128 local expert base weights. On B200, the pinned Transformer Engine grouped path launched roughly one small GEMM per expert, leaving substantial launch overhead. The opt-in torch backend packs the frozen expert weights once and runs the base branch through torch.nn.functional.grouped_mm when available, with torch._grouped_mm as the compatibility fallback. The ordinary Bridge LoRA branch remains unchanged.

This path reduced the measured logprob-plus-update cycle from 21.640 s to 14.684 s on the production-shaped benchmark: 32.1% less time and 47.4% more throughput, with the same peak allocated memory.

B200 throughput evidence

The comparison used the same node, step-75 checkpoint, fixed production-mixed token batch, and TP4/PP1/EP4/ETP1 topology. Both arms used BF16 LoRA, ordinary all-to-all dispatch, a 32k dynamic token cap, and 121,088 input tokens per iteration.

Artifacts:

• TE reference: job 102251, tf-super-deepep-ab-v1-a4e9b214c0/super-topology-benchmark.json, tp4-alltoall arm.
• Torch candidate: job 102265, tf-super-torch-gmm-v1-a518a324d7/super-topology-benchmark.json.
• Kernel profile: job 102277, tf-super-torch-gmm-profile-v1-a518a324d7.

The parity gate reported zero relative-L2 and max-absolute error for both forward output and input gradient, including experts with zero assigned tokens. The runtime packed 80 grouped FC1/FC2 modules and relocated 52.5 GiB of frozen weights without adding checkpoint entries.

Why existing behavior remains safe

• transformer_engine remains the default. No model uses the new path unless its config explicitly selects moe_expert_gemm_backend="torch".
• The torch path accepts only grouped, frozen, bias-free BF16 expert weights. FP8, FP4, trainable base weights, non-MoE models, TE single-grouped weights, the TE operation fuser, delayed weight-gradient modes, local transformer specs, and TE builds without grouped-linear support fail during configuration or preparation.
• Existing Parameter objects and state-dict keys are preserved. Each expert parameter becomes a view into one non-persistent contiguous buffer.
• Expert split metadata is checked before both preparation and the empty-input fast path.
• Public grouped GEMM is preferred; the private PyTorch operation remains a fallback for older supported images.
• Tests cover TE-versus-torch output and input-gradient parity, empty batches, malformed and zero-token expert splits, public-only and private-only API environments, checkpoint-schema preservation, and recovery after a middle expert weight is replaced.
• Downstream NeMo-RL tests compare ordinary Bridge LoRA over TE and torch base GEMMs, including output, input-gradient, and adapter-gradient parity.

Full-recompute correctness

The PR also carries the small frozen-input guard needed by PP1 adapter-only full recompute. Re-entrant checkpointing otherwise sees a frozen embedding output with requires_grad=False and can detach every checkpointed chunk from the adapters inside it.

The regression freezes the block and input, installs trainable low-rank adapters, runs backward, and requires nonzero gradients on every adapter parameter. It covers uniform recompute and block recompute; the block case exercises both a checkpointed layer and the uncheckpointed remainder. The production 32k and 48k NeMo-RL runs use this full-recompute path; their recorded gradients are nonzero and their losses move.

Scope deliberately omitted

• Fused expert LoRA was measured at 53.137 s versus 53.192 s for ordinary LoRA and used about 1.16 GiB more peak memory, so it is not included.
• Selective expert_fc1_act recompute was faster than whole-MoE recompute at short context, but the 32k selective arm OOMed because it did not cover the 40 Mamba mixers. The tested 32k and 48k recipes use full recompute instead.

Validation status

• Black, isort, Ruff, Python compilation, and git diff --check pass at final head c03fbc4a1.
• The required fork-base check passes, every PR commit carries DCO sign-off, Sam's review is approved, and all five follow-up review threads are addressed and resolved.
• Downstream NeMo-RL PR #60 passed its pre-commit, lockfile, fork-base, build-deploy, and GPU-gate checks with the implementation head 8447c0d3.
• End-to-end B200 validation of 8447c0d3 is provided by NeMo-RL job 113785: it used the torch backend, PP1 full recompute, and ordinary LoRA, migrated the step-75 adapter checkpoint, and completed optimizer steps with nonzero gradient norms and changing rewards in W&B run as6tqeze.
• Final head c03fbc4a1 adds review hardening around unsupported configuration combinations, public/private API selection, malformed split validation, and broader regression coverage. A separate one-B200 targeted run, job 170738, is queued against that exact commit; it does not replace or share GPUs with the production run.

Commit signing

All commits carry DCO sign-off. Cryptographic signing was unavailable because this machine has no configured GPG or SSH signing identity.